# =============================================================================
# Statistical Analysis Code for:
# "Chronic adaptive versus conventional DBS response patterns in Parkinson's
#  disease: A pilot randomized crossover trial"
#
# Visualization: Forest plots for effect sizes (Cohen's d) with 95% CI
# Author: Jun Tanimura, MD, MSc.
# Created: 2025-08-07
# Last Updated: 2025-11-26
# Version: v3
#
# CHANGES FROM v2 FOR PIEER REVIEW REVISION:
# - Changed from lollipop plot to forest plot style (point + error bar)
# - Added 95% CI for Cohen's d effect sizes using Effect_Size_CI_Lower/Upper
# =============================================================================

# Load required packages
library(tidyverse)
library(ggplot2)
library(patchwork)

# Avoid conflicts by explicitly using scales functions
require(scales)

# Read the FDR results
results_df <- read_csv("results_standard_crossover.csv")

# Verify that Cohen's d CI columns exist (required for v3 forest plot)
required_cols <- c("Effect_Size", "Effect_Size_CI_Lower", "Effect_Size_CI_Upper")
if(!all(required_cols %in% names(results_df))) {
  stop("ERROR: The CSV file is missing Cohen's d CI columns (Effect_Size_CI_Lower, Effect_Size_CI_Upper).\nPlease run 02_ANCOVA_v3.1_FULL.R to generate the updated results file.")
}

# Define clinical evaluation categories and clean names (same order as reference code)
evaluation_categories <- tribble(
  ~Evaluation, ~Category, ~Clean_Name, ~Order,
  "Mean_corr_ON_time", "Motor Diary", "ON Time", 1,
  "Mean_corr_OFF_time", "Motor Diary", "OFF Time", 2,
  "Mean_corr_Dys_sev_time", "Motor Diary", "Troublesome Dyskinesia Time", 3,
  "Mean_corr_Dys_weak_time", "Motor Diary", "Non-troublesome Dyskinesia Time", 4,
  "ADL", "Functional", "ADL", 5,
  "UPDRS_part_1", "UPDRS", "UPDRS Part I", 6,
  "UPDRS_part_2", "UPDRS", "UPDRS Part II", 7,
  "UPDRS_part_3", "UPDRS", "UPDRS Part III", 8,
  "UPDRS_part_4", "UPDRS", "UPDRS Part IV", 9,
  "UPDRS_total", "UPDRS", "UPDRS Total Score", 10,
  "MMSE", "Cognitive", "MMSE", 11,
  "PDQ-39", "QOL", "PDQ 39", 12
)

# Merge with categories and prepare data
plot_data <- results_df %>%
  left_join(evaluation_categories, by = "Evaluation") %>%
  filter(!is.na(Clean_Name)) %>%
  mutate(
    # Create significance categories
    Significance_Status = case_when(
      Significant_FDR ~ "Significant (q < 0.05)",
      Significant_uncorrected ~ "Lost significance after FDR",
      Treatment_P_FDR < 0.10 ~ "Marginal (q < 0.10)",
      TRUE ~ "Non-significant"
    ),
    Significance_Status = factor(Significance_Status, levels = c(
      "Significant (q < 0.05)", 
      "Lost significance after FDR",
      "Marginal (q < 0.10)", 
      "Non-significant"
    )),
    
    # Effect size categories
    Effect_Category = case_when(
      abs(Effect_Size) >= 0.8 ~ "Large (|d| ≥ 0.8)",
      abs(Effect_Size) >= 0.5 ~ "Medium (|d| ≥ 0.5)",
      abs(Effect_Size) >= 0.2 ~ "Small (|d| ≥ 0.2)",
      TRUE ~ "Negligible (|d| < 0.2)"
    ),
    Effect_Category = factor(Effect_Category, levels = c(
      "Large (|d| ≥ 0.8)", "Medium (|d| ≥ 0.5)", 
      "Small (|d| ≥ 0.2)", "Negligible (|d| < 0.2)"
    )),
    
    # Direction of effect
    Effect_Direction = ifelse(Treatment_Estimate < 0, "Favors aDBS", "Favors cDBS"),
    
    # Order by the predefined sequence
    Clean_Name = factor(Clean_Name, levels = evaluation_categories$Clean_Name[evaluation_categories$Order])
  ) %>%
  arrange(Order)

# Define color schemes
effect_colors <- c(
  "Significant (q < 0.05)" = "#D32F2F",
  "Lost significance after FDR" = "#FF9800", 
  "Marginal (q < 0.10)" = "#FFC107",
  "Non-significant" = "#757575"
)

# Effect size category colors (for supplement figure)
effect_category_colors <- c(
  "Large (|d| ≥ 0.8)" = "#D32F2F",
  "Medium (|d| ≥ 0.5)" = "#FF9800",
  "Small (|d| ≥ 0.2)" = "#FFC107",
  "Negligible (|d| < 0.2)" = "#757575"
)

# Simple direction-based colors
direction_colors <- c(
  "aDBS_Superior" = "darkred",
  "cDBS_Superior" = "steelblue",
  "Favors aDBS" = "darkred",
  "Favors cDBS" = "steelblue"
)

# ===========================================
# MAIN FIGURE: Faceted Forest Plot with Effect Sizes
# ===========================================

create_simple_cohens_d_plot <- function(data) {
  
  # Prepare data with consistent ordering
  plot_data_ordered <- data %>%
    arrange(Order) %>%
    mutate(
      # Simple direction classification
      Direction_Simple = ifelse(Effect_Size > 0, "aDBS_Superior", "cDBS_Superior"),
      # Order for plotting (bottom to top)
      Clean_Name = factor(Clean_Name, levels = rev(Clean_Name))
    )
  
  # Calculate dynamic x-axis limits based on data INCLUDING 95% CI
  min_value <- min(c(plot_data_ordered$Effect_Size, plot_data_ordered$Effect_Size_CI_Lower), na.rm = TRUE)
  max_value <- max(c(plot_data_ordered$Effect_Size, plot_data_ordered$Effect_Size_CI_Upper), na.rm = TRUE)
  x_range <- max_value - min_value
  x_limits <- c(min_value - x_range * 0.15, max_value + x_range * 0.15)
  
  # Create the Cohen's D plot with 95% CI (forest plot style)
  p_cohens_d <- ggplot(plot_data_ordered, aes(y = Clean_Name)) +
    
    # Add vertical line at zero
    geom_vline(xintercept = 0, linetype = "solid", color = "gray50", linewidth = 0.5) +
    
    # Add simple dotted lines for effect size thresholds (only if within range)
    {if(x_limits[2] >= 0.2) geom_vline(xintercept = c(-0.2, 0.2), linetype = "dotted", color = "gray70", alpha = 0.7)} +
    {if(x_limits[2] >= 0.5) geom_vline(xintercept = c(-0.5, 0.5), linetype = "dotted", color = "gray70", alpha = 0.7)} +
    {if(x_limits[2] >= 0.8) geom_vline(xintercept = c(-0.8, 0.8), linetype = "dotted", color = "gray70", alpha = 0.7)} +
    
    # Add error bars for 95% CI (horizontal bars)
    geom_errorbarh(aes(xmin = Effect_Size_CI_Lower, xmax = Effect_Size_CI_Upper, 
                       color = Direction_Simple),
                   height = 0.3, linewidth = 0.8, alpha = 0.7) +
    
    # Add points for effect sizes
    geom_point(aes(x = Effect_Size, color = Direction_Simple), 
               size = 3.5, alpha = 0.9) +
    
    # Add text labels showing exact values
    geom_text(aes(x = Effect_Size, label = sprintf("%.2f", Effect_Size)),
              hjust = -0.3, vjust = 0.5, size = 3, color = "black") +
    
    # Color scale
    scale_color_manual(
      values = direction_colors,
      guide = "none"  # Remove legend
    ) +
    
    # Dynamic X-axis
    scale_x_continuous(
      limits = x_limits,
      breaks = scales::pretty_breaks(n = 6),
      labels = function(x) sprintf("%.1f", x)
    ) +
    
    # Labels
    labs(
      title = "Effect Sizes with 95% Confidence Intervals: aDBS vs cDBS",
      x = "Cohen's d (95% CI)",
      y = "",
      caption = "Positive values indicate aDBS superiority\nError bars represent 95% confidence intervals\nFDR correction applied using Benjamini-Hochberg method"
    ) +
    
    # Theme
    theme_classic(base_size = 12) +
    theme(
      plot.title = element_text(size = 14, face = "bold"),
      axis.text.y = element_text(size = 10),
      axis.text.x = element_text(size = 9),
      axis.title.x = element_text(size = 10),
      panel.grid.major.x = element_blank(),  # Remove grid lines
      plot.caption = element_text(size = 8, color = "gray50", hjust = 0),
      plot.margin = margin(10, 10, 10, 10)
    )
  
  # Add threshold annotations only if thresholds are visible
  if(x_limits[2] >= 0.2) {
    p_cohens_d <- p_cohens_d + 
      annotate("text", x = 0.2, y = 1, label = "Small", 
               angle = 90, vjust = -0.5, size = 3, color = "gray60")
  }
  
  if(x_limits[2] >= 0.5) {
    p_cohens_d <- p_cohens_d + 
      annotate("text", x = 0.5, y = 1, label = "Medium", 
               angle = 90, vjust = -0.5, size = 3, color = "gray60")
  }
  
  if(x_limits[2] >= 0.8) {
    p_cohens_d <- p_cohens_d + 
      annotate("text", x = 0.8, y = 1, label = "Large", 
               angle = 90, vjust = -0.5, size = 3, color = "gray60")
  }
  
  return(p_cohens_d)
}

# ===========================================
# SUPPLEMENT FIGURE: Effect Size Summary
# ===========================================

create_effect_size_summary <- function(data) {
  
  # Prepare effect size data
  effect_data <- data %>%
    mutate(
      Clean_Name = fct_reorder(Clean_Name, Order),
      Effect_Direction = ifelse(Effect_Size < 0, "Favors aDBS", "Favors cDBS")
    )
  
  p_effect <- ggplot(effect_data, aes(y = Clean_Name, x = abs(Effect_Size))) +
    
    # Add vertical lines for effect size thresholds
    geom_vline(xintercept = 0.2, linetype = "dotted", color = "gray60", alpha = 0.7) +
    geom_vline(xintercept = 0.5, linetype = "dashed", color = "gray60", alpha = 0.7) +
    geom_vline(xintercept = 0.8, linetype = "solid", color = "gray60", alpha = 0.7) +
    
    # Add bars
    geom_col(aes(fill = Effect_Category), alpha = 0.8, width = 0.7) +
    
    # Add text labels for effect sizes
    geom_text(aes(label = sprintf("%.2f", Effect_Size),
                  x = abs(Effect_Size) + 0.05),
              hjust = 0, size = 3) +
    
    # Add direction indicators
    geom_point(aes(color = Effect_Direction), 
               x = 0, size = 3, alpha = 0.8) +
    
    # Customization
    scale_fill_manual(values = effect_category_colors, name = "Effect Size") +
    scale_color_manual(values = direction_colors, name = "Direction") +
    
    labs(
      title = "Effect Sizes (Cohen's d) by Outcome Measure",
      x = "Absolute Effect Size |d|",
      y = "",
      caption = "Vertical lines: Small (0.2), Medium (0.5), Large (0.8) effect thresholds"
    ) +
    
    theme_classic(base_size = 11) +
    theme(
      plot.title = element_text(size = 12, face = "bold"),
      legend.position = "right",
      legend.text = element_text(size = 9),
      plot.caption = element_text(size = 8, color = "gray50")
    ) +
    
    # Add threshold annotations
    annotate("text", x = 0.2, y = 1, label = "Small", 
             angle = 90, vjust = -0.5, size = 3, color = "gray60") +
    annotate("text", x = 0.5, y = 1, label = "Medium", 
             angle = 90, vjust = -0.5, size = 3, color = "gray60") +
    annotate("text", x = 0.8, y = 1, label = "Large", 
             angle = 90, vjust = -0.5, size = 3, color = "gray60")
  
  return(p_effect)
}

# ===========================================
# CREATE AND SAVE PLOTS
# ===========================================

# Generate plots
main_plot <- create_simple_cohens_d_plot(plot_data)
effect_plot <- create_effect_size_summary(plot_data)

# Create supplement figure
supplement_figure <- effect_plot + 
  plot_annotation(
    title = "Supplementary Analysis: Effect Sizes by Outcome Measure",
    theme = theme(plot.title = element_text(size = 14, face = "bold"))
  )

# Save main figure (PDF only) - forest plot with 95% CI
ggsave("figure_forest_plot_cohens_d.pdf", 
       main_plot, width = 4.4, height = 3.4, bg = "white")

# Save supplement figure (PDF only)  
ggsave("figure_supplement_effect_sizes.pdf", 
       supplement_figure, width = 8, height = 3.4, bg = "white")

# Print summary
cat("\n=== VISUALIZATION COMPLETE ===\n")
cat("Files created:\n")
cat("1. Main_Figure_Forest_Plot_Cohens_D_with_CI.pdf - Forest plot with 95% CI\n")
cat("2. Supplement_Figure_Effect_Sizes.pdf - Effect sizes ranking\n")

# Check data integrity
cat("\n=== DATA INTEGRITY CHECK ===\n")

# Check treatment estimate CI
data_check <- plot_data %>%
  select(Clean_Name, Treatment_Estimate, CI_Lower, CI_Upper) %>%
  mutate(
    CI_Contains_Estimate = (Treatment_Estimate >= CI_Lower & Treatment_Estimate <= CI_Upper),
    CI_Width = CI_Upper - CI_Lower
  )

cat("Checking if Treatment_Estimate falls within CI_Lower and CI_Upper:\n")
print(data_check)

problematic_data <- data_check %>% filter(!CI_Contains_Estimate)
if(nrow(problematic_data) > 0) {
  cat("\n*** WARNING: Found problematic treatment estimate confidence intervals ***\n")
  print(problematic_data)
} else {
  cat("\nAll treatment estimate confidence intervals appear correct.\n")
}

# Check Cohen's d CI
cohens_d_check <- plot_data %>%
  select(Clean_Name, Effect_Size, Effect_Size_CI_Lower, Effect_Size_CI_Upper) %>%
  mutate(
    CI_Contains_Effect = (Effect_Size >= Effect_Size_CI_Lower & Effect_Size <= Effect_Size_CI_Upper),
    CI_Width = Effect_Size_CI_Upper - Effect_Size_CI_Lower
  )

cat("\nChecking if Effect_Size falls within Effect_Size_CI_Lower and Effect_Size_CI_Upper:\n")
print(cohens_d_check)

problematic_cohens_d <- cohens_d_check %>% filter(!CI_Contains_Effect)
if(nrow(problematic_cohens_d) > 0) {
  cat("\n*** WARNING: Found problematic Cohen's d confidence intervals ***\n")
  print(problematic_cohens_d)
} else {
  cat("\nAll Cohen's d confidence intervals appear correct.\n")
}

# Display main plot
print(main_plot)
cat("\nForest plot with 95% CI created successfully!\n")

# Clean up to avoid conflicts
detach("package:scales", unload = TRUE)